package tasks

import (
	"fmt"
	"mediumkube/pkg/common"
	"mediumkube/pkg/network"
	"mediumkube/pkg/utils"
)

// IPMode how rules are inserted into iptables

const (
	table string = "filter"
)

var (
	ruleRegistry [][]string = make([][]string, 0)
)

func _forwardRuleIn(bridge common.Bridge) []string {
	return []string{
		"-s", bridge.Inet,
		"-i", bridge.Name,
		"-j", "ACCEPT",
	}
}

func _forwardRuleOut(bridge common.Bridge) []string {
	return []string{
		"-d", bridge.Inet,
		"-o", bridge.Name,
		"-m", "conntrack",
		"--ctstate", "RELATED,ESTABLISHED",
		"-j", "ACCEPT",
	}
}

func _forwardRuleIO(bridge common.Bridge) []string {
	return []string{
		"-i", bridge.Name,
		"-o", bridge.Name,
		"-j", "ACCEPT",
	}
}

func _forwardRejectICMPUnreachableIn(bridge common.Bridge) []string {
	return []string{
		"-i", bridge.Name,
		"-j", "REJECT",
		"--reject-with", "icmp-port-unreachable",
	}
}

func _forwardRejectICMPUnreachableOut(bridge common.Bridge) []string {
	return []string{
		"-o", bridge.Name,
		"-j", "REJECT",
		"--reject-with", "icmp-port-unreachable",
	}
}

func _dhcpIn(bridge common.Bridge) []string {
	return []string{
		"-i", bridge.Name,
		"-p", "udp",
		"-m", "udp",
		"--dport", "67",
		"-j", "ACCEPT",
	}
}

func _dhcpOut(bridge common.Bridge) []string {
	return []string{
		"-o", bridge.Name,
		"-p", "udp",
		"-m", "udp",
		"--sport", "67",
		"-j", "ACCEPT",
	}
}

func _dnsIn(bridge common.Bridge, protocol string) []string {
	return []string{
		"-i", bridge.Name,
		"-p", protocol,
		"-m", protocol,
		"--dport", "53",
		"-j", "ACCEPT",
	}
}

func _dnsOut(bridge common.Bridge, protocol string) []string {
	return []string{
		"-o", bridge.Name,
		"-p", protocol,
		"-m", protocol,
		"--dport", "53",
		"-j", "ACCEPT",
	}
}

func _in(port int) []string {
	return []string{
		"-p", "tcp",
		"--dport", fmt.Sprintf("%v", port),
		"-j", "ACCEPT",
	}
}

func _out(port int) []string {
	return []string{
		"-p", "tcp",
		"--sport", fmt.Sprintf("%v", port),
		"-j", "ACCEPT",
	}
}

func insertRuleIfNotExists(chain string, mode network.IPMode, rules ...string) {
	rules = append(rules, "-m", "comment", "--comment", "Auto generated by mediumkubed")
	if !utils.ContainsSlice(ruleRegistry, append([]string{chain}, rules...)) {
		ruleRegistry = append(ruleRegistry, append([]string{chain}, rules...))
	}
	network.InsertRuleIfNotExists(network.TableFilter, chain, mode, rules...)
}

// ProcessIptables configure netfilter rules required for DNS, DHCP and other applications
func ProcessIptables(config *common.OverallConfig) {
	bridge := config.Bridge
	insertRuleIfNotExists("FORWARD", network.IPModPREP, _forwardRuleOut(bridge)...) // Allow outbound traffic from bridge
	insertRuleIfNotExists("FORWARD", network.IPModPREP, _forwardRuleIn(bridge)...)  // Allow inbound traffic to bridge
	insertRuleIfNotExists("FORWARD", network.IPModPREP, _forwardRuleIO(bridge)...)
	insertRuleIfNotExists("FORWARD", network.IPModAPP, _forwardRejectICMPUnreachableIn(bridge)...) // Reject traffic when ICMP unreachable
	insertRuleIfNotExists("FORWARD", network.IPModAPP, _forwardRejectICMPUnreachableOut(bridge)...)
	insertRuleIfNotExists("INPUT", network.IPModPREP, _dhcpIn(bridge)...) // Open port for DHCP
	insertRuleIfNotExists("INPUT", network.IPModPREP, _dnsIn(bridge, "tcp")...)
	insertRuleIfNotExists("INPUT", network.IPModPREP, _dnsIn(bridge, "udp")...)
	insertRuleIfNotExists("OUTPUT", network.IPModPREP, _dhcpOut(bridge)...)
	insertRuleIfNotExists("OUTPUT", network.IPModPREP, _dnsOut(bridge, "tcp")...)
	insertRuleIfNotExists("OUTPUT", network.IPModPREP, _dnsOut(bridge, "udp")...)
	insertRuleIfNotExists("INPUT", network.IPModAPP, _in(config.Overlay.EtcdPort)...)
	insertRuleIfNotExists("INPUT", network.IPModAPP, _in(config.Overlay.GRPCPort)...)
}

// CleanUpIptables clean up iptables when exiting
func CleanUpIptables() {

	for _, cr := range ruleRegistry {
		chain := cr[0]
		rules := cr[1:]
		network.DeleteIfExists(network.TableFilter, chain, rules)
	}
}
